#include <ros/ros.h>
#include <sensor_msgs/PointCloud2.h>
#include <tf/transform_listener.h>
#include <pcl_ros/transforms.h>
#include <pcl_conversions/pcl_conversions.h>
#include <pcl/point_types.h>
#include <pcl/filters/extract_indices.h>
#include <pcl/sample_consensus/model_types.h>
#include <pcl/sample_consensus/method_types.h>
#include <pcl/segmentation/sac_segmentation.h>
#include <pcl/filters/passthrough.h>
#include <pcl/filters/voxel_grid.h>
#include <chrono>
#include <nav_msgs/OccupancyGrid.h>

class GroundSegmentationNode {
public:
    GroundSegmentationNode() : leaf_size_(0.02f), last_adjust_time_(ros::Time::now()) {
        ROS_INFO("Initializing GroundSegmentationNode...");
        ros::NodeHandle nh;
        ros::NodeHandle pnh("~");
        std::string input_topic;
        pnh.param<std::string>("input_topic", input_topic, "/camera/depth/color/points");
        cloud_sub_ = nh.subscribe(input_topic, 1, &GroundSegmentationNode::cloudCallback, this);
        ground_pub_ = nh.advertise<sensor_msgs::PointCloud2>("/ground_points", 1);
        obstacle_pub_ = nh.advertise<sensor_msgs::PointCloud2>("/obstacle_points", 1);
        grid_pub_ = nh.advertise<nav_msgs::OccupancyGrid>("/obstacle_grid", 1);
        ROS_INFO("Finished initializing GroundSegmentationNode.");
    }

    void cloudCallback(const sensor_msgs::PointCloud2ConstPtr& msg) {
        auto t1 = std::chrono::high_resolution_clock::now();
        // Transform to base_link
        sensor_msgs::PointCloud2 cloud_base;
        try {
            listener_.waitForTransform("base_link", msg->header.frame_id, msg->header.stamp, ros::Duration(0.2));
            pcl_ros::transformPointCloud("base_link", *msg, cloud_base, listener_);
        } catch (tf::TransformException& ex) {
            ROS_WARN("TF transform failed: %s", ex.what());
            return;
        }
        // Convert to PCL
        pcl::PointCloud<pcl::PointXYZ>::Ptr cloud(new pcl::PointCloud<pcl::PointXYZ>);
        pcl::fromROSMsg(cloud_base, *cloud);
        // ROI裁剪
        pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_roi(new pcl::PointCloud<pcl::PointXYZ>);
        pcl::PassThrough<pcl::PointXYZ> pass;
        pass.setInputCloud(cloud);
        pass.setFilterFieldName("x"); pass.setFilterLimits(0.0, 6.0); pass.filter(*cloud_roi);
        pass.setInputCloud(cloud_roi);
        pass.setFilterFieldName("y"); pass.setFilterLimits(-5.0, 5.0); pass.filter(*cloud_roi);
        pass.setInputCloud(cloud_roi);
        pass.setFilterFieldName("z"); pass.setFilterLimits(-1.0, 2.0); pass.filter(*cloud_roi);
        // VoxelGrid降采样
        pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_ds(new pcl::PointCloud<pcl::PointXYZ>);
        pcl::VoxelGrid<pcl::PointXYZ> voxel;
        voxel.setInputCloud(cloud_roi);
        voxel.setLeafSize(leaf_size_, leaf_size_, leaf_size_);
        voxel.filter(*cloud_ds);
        // RANSAC plane segmentation
        pcl::SACSegmentation<pcl::PointXYZ> seg;
        seg.setOptimizeCoefficients(true);
        seg.setModelType(pcl::SACMODEL_PLANE);
        seg.setMethodType(pcl::SAC_RANSAC);
        seg.setDistanceThreshold(0.03);
        seg.setInputCloud(cloud_ds);
        pcl::ModelCoefficients::Ptr coeffs(new pcl::ModelCoefficients);
        pcl::PointIndices::Ptr inliers(new pcl::PointIndices);
        seg.segment(*inliers, *coeffs);
        if (inliers->indices.empty()) {
            ROS_WARN("No ground plane found.");
            return;
        }
        // Extract ground
        pcl::ExtractIndices<pcl::PointXYZ> extract;
        extract.setInputCloud(cloud_ds);
        extract.setIndices(inliers);
        extract.setNegative(false);
        pcl::PointCloud<pcl::PointXYZ>::Ptr ground(new pcl::PointCloud<pcl::PointXYZ>);
        extract.filter(*ground);
        // Extract obstacles
        extract.setNegative(true);
        pcl::PointCloud<pcl::PointXYZ>::Ptr obstacles_all(new pcl::PointCloud<pcl::PointXYZ>);
        extract.filter(*obstacles_all);
        // 进一步去除低于地面的点
        pcl::PointCloud<pcl::PointXYZ>::Ptr obstacles(new pcl::PointCloud<pcl::PointXYZ>);
        float a = coeffs->values[0], b = coeffs->values[1], c = coeffs->values[2], d = coeffs->values[3];
        float min_dist = 0.15;
        for (const auto& pt : obstacles_all->points) {
            float dist = a * pt.x + b * pt.y + c * pt.z + d;
            if (dist > min_dist) obstacles->points.push_back(pt);
        }
        obstacles->width = obstacles->points.size();
        obstacles->height = 1;
        obstacles->is_dense = false;
        // Publish
        sensor_msgs::PointCloud2 ground_msg, obstacle_msg;
        pcl::toROSMsg(*ground, ground_msg);
        pcl::toROSMsg(*obstacles, obstacle_msg);
        ground_msg.header = cloud_base.header;
        obstacle_msg.header = cloud_base.header;
        ground_pub_.publish(ground_msg);
        obstacle_pub_.publish(obstacle_msg);
        // OccupancyGrid
        nav_msgs::OccupancyGrid grid;
        grid.header = cloud_base.header;
        grid.info.resolution = 0.1;
        grid.info.width = 100;
        grid.info.height = 100;
        grid.info.origin.position.x = -5.0;
        grid.info.origin.position.y = -5.0;
        grid.info.origin.position.z = 0.0;
        grid.info.origin.orientation.w = 1.0;
        grid.data.assign(100 * 100, 0); // 0: free, 100: occupied
        for (const auto& pt : obstacles->points) {
            int ix = static_cast<int>((pt.x + 5.0) / 0.1);
            int iy = static_cast<int>((pt.y + 5.0) / 0.1);
            if (ix >= 0 && ix < 100 && iy >= 0 && iy < 100) {
                int idx = iy * 100 + ix;
                grid.data[idx] = 100;
            }
        }
        grid_pub_.publish(grid);
        // 统计耗时
        auto t2 = std::chrono::high_resolution_clock::now();
        double elapsed_ms = std::chrono::duration<double, std::milli>(t2 - t1).count();
        ROS_INFO("Frame processed in %.2f ms, input: %zu, ROI: %zu, downsampled: %zu", elapsed_ms, cloud->size(), cloud_roi->size(), cloud_ds->size());
        ROS_INFO("Leaf size: %.3f", leaf_size_);
        // 动态调节LeafSize
        ros::Time now = ros::Time::now();
        if ((now - last_adjust_time_).toSec() > 1.0) {
            if (elapsed_ms > 100.0 && leaf_size_ < 0.5f) {
                leaf_size_ *= 1.2f;
                ROS_INFO("[VoxelGrid] Increase leaf_size to %.3f", leaf_size_);
            } else if (elapsed_ms < 80.0 && leaf_size_ > 0.01f) {
                leaf_size_ /= 1.2f;
                ROS_INFO("[VoxelGrid] Decrease leaf_size to %.3f", leaf_size_);
            }
            last_adjust_time_ = now;
        }
    }

private:
    ros::Subscriber cloud_sub_;
    ros::Publisher ground_pub_;
    ros::Publisher obstacle_pub_;
    ros::Publisher grid_pub_;
    tf::TransformListener listener_;
    float leaf_size_;
    ros::Time last_adjust_time_;
};

int main(int argc, char** argv) {
    ros::init(argc, argv, "ground_segmentation_node");
    std::cout << "Starting ground segmentation node..." << std::endl;
    GroundSegmentationNode node;
    ros::spin();
    return 0;
}
